import base64
import json
import os
import os.path as osp
import labelme
import numpy as np
from labelme import utils
from PIL import Image

if __name__ == '__main__':

    # 原图文件夹
    jpgs_path = r"F:\IAST_SEG_Test\image"
    os.makedirs(jpgs_path, exist_ok=True)

    # 代码执行完生成的color文件夹
    jpgscolor_path = r"F:\IAST_SEG_Test\imagecolor"
    os.makedirs(jpgscolor_path, exist_ok=True)

    # 代码执行完生成的mask掩码文件夹
    pngs_path = r"F:\IAST_SEG_Test\imagePNG"
    os.makedirs(pngs_path, exist_ok=True)

    # 标注生成的json文件夹
    json_path = r"F:\IAST_SEG_Test\imagejson"

    classes = ["__background__", "grass", "road", "farmland", "river", "lake", "container", "building", "forest",
               "bare-soil", "paved-surface", "railway", "greenhouse", "stadium"]
    label_name_to_value = dict(zip(classes, range(len(classes))))
    colormap = np.array([
        [0, 0, 0],  # 0 background
        [8, 248, 0],  # 1 grass
        [208, 244, 247],  # 2 road
        [193, 95, 96],  # 3 farmland
        [89, 164, 255],  # 4 river
        [42, 14, 255],  # 5 lake
        [148, 82, 255],  # 6 container
        [255, 65, 220],  # 7 building
        [8, 127, 0],  # 8 forest
        [103, 27, 4],  # 9 bare-soil
        [118, 128, 164],  # 10 paved-surface
        [255, 108, 16],  # 11 railway
        [140, 41, 133],  # 12 greenhouse
        [255, 88, 98],  # 13 stadium
    ], dtype=np.uint8)

    json_files = os.listdir(json_path)
    for json_file in json_files:
        path = os.path.join(json_path, json_file)

        if os.path.isfile(path) and path.endswith('json'):
            data = json.load(open(path))

            # if data['imageData']:
            #     imageData = data['imageData']
            # else:
            #     # imagePath = os.path.join(os.path.dirname(path), data['imagePath'])
            #     imagePath = os.path.join(data['info']['folder'],data['info']['name'])
            #     with open(imagePath, 'rb') as f:
            #         imageData = f.read()
            #         imageData = base64.b64encode(imageData).decode('utf-8')

            imagePath = os.path.join(data['info']['folder'], data['info']['name'])
            with open(imagePath, 'rb') as f:
                imageData = f.read()
                imageData = base64.b64encode(imageData).decode('utf-8')
            img = utils.img_b64_to_arr(imageData)

            lbl = utils.shapes_to_label(img.shape, data['objects'], label_name_to_value)
            # 将标签映射到类索引
            new = np.zeros(lbl.shape, dtype=int)
            for label_name, label_value in label_name_to_value.items():
                new[lbl == label_value] = classes.index(label_name)

            # 应用颜色映射
            color_label_img = colormap[new]

            Image.fromarray(img).save(osp.join(jpgs_path, json_file.split(".")[0] + '.jpg'))
            Image.fromarray(color_label_img).save(osp.join(jpgscolor_path, json_file.split(".")[0] + '_color.png'))

            utils.lblsave(osp.join(pngs_path, json_file.split(".")[0] + '.png'), lbl)
            print('Saved ' + json_file.split(".")[0] + '.jpg, ' + json_file.split(".")[0] + '.png, and ' +
                  json_file.split(".")[0] + '_color.png')
